import time
import numpy as np  # 用于处理多维数组和矩阵运算
import faiss  # 用于高效相似性搜索和稠密向量聚类
from xgboost.testing.data import joblib
from util import createXY  # 用于创建数据集的特征和标签
from sklearn.model_selection import train_test_split  # 用于拆分数据集为训练集和测试集
import argparse  # 用于解析命令行参数
import logging  # 用于记录日志
import pickle
import gradio as gr
import cv2  # 导入cv2模块,用于读取和处理图像
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier, VotingClassifier, BaggingClassifier, AdaBoostClassifier
from sklearn.svm import SVC
from sklearn.tree import DecisionTreeClassifier
from xgboost import XGBClassifier
from sklearn.ensemble import StackingClassifier
from tabulate import tabulate

# 配置logging, 确保能够打印正在运行的函数名
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# 获取命令行参数
def get_args():
    args = argparse.Namespace()
    args.mode = 'cpu'
    args.feature = 'flat'
    args.library = 'faiss'
    return args

def main():
    args = get_args()

    # 初始化FAISS所需的资源
    res = None
    logging.info(f"选择模式是 {args.mode.upper()}")
    logging.info(f"选择特征提取方法是 {args.feature.upper()}")
    logging.info(f"选择使用的库是 {args.library.upper()}")

    # 载入和预处理数据
    X, y = createXY(train_folder="data/train", dest_folder=".", method=args.feature)
    X = np.array(X).astype('float32')
    faiss.normalize_L2(X)  # 对数据进行L2归一化
    y = np.array(y)
    logging.info("数据加载和预处理完成。")

    # 数据集分割为训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=2023)
    logging.info("数据集划分为训练集和测试集。")

    # 创建逻辑回归分类器
    logistic_clf = LogisticRegression(max_iter=1000)

    # 创建随机森林分类器
    random_forest_clf = RandomForestClassifier(random_state=42)

    # 创建支持向量机分类器
    # probability=True表示需要计算概率值，这样才能使用软投票
    svm_clf = SVC(probability=True)

    # 创建硬投票分类器，包含逻辑回归、随机森林和支持向量机三个分类器
    voting_clf_hard = VotingClassifier(
        estimators=[('lr', logistic_clf), ('rf', random_forest_clf), ('svc', svm_clf)],
        voting='hard'  # 采用硬投票, 即选择票数最多的类别作为预测结果
    )

    # 创建软投票分类器，包含逻辑回归、随机森林和支持向量机三个分类器
    voting_clf_soft = VotingClassifier(
        estimators=[('lr', logistic_clf), ('rf', random_forest_clf), ('svc', svm_clf)],
        voting='soft'  # 采用软投票, 即选择概率之和最大的类别作为预测结果
    )

    # 创建Bagging分类器
    bag_clf = BaggingClassifier(
        DecisionTreeClassifier(),  # 基分类器选用决策树分类器
        n_estimators=500,  # 500个分类器
        max_samples=100,  # 每个分类器的训练集包含100个样本
        bootstrap=True,  # 有放回的采样
        n_jobs=-1,  # 使用所有CPU核
    )

    # 创建Pasting分类器
    paste_clf = BaggingClassifier(
        DecisionTreeClassifier(),  # 基分类器选用决策树分类器
        n_estimators=500,  # 500个分类器
        max_samples=100,  # 每个分类器的训练集包含100个样本
        bootstrap=False,  # 无放回的采样
        n_jobs=-1,  # 使用所有CPU核
    )

    # 创建AdaBoost分类器
    ada_clf = AdaBoostClassifier(
        DecisionTreeClassifier(max_depth=1),  # 基分类器选用决策树分类器
        n_estimators=200,  # 200个分类器
        algorithm="SAMME.R",  # 使用SAMME.R算法
        learning_rate=0.5  # 学习率为0.5, 即每个分类器的权重缩减系数为0.5
    )

    # 创建一个梯度提升分类器
    xgb_clf = XGBClassifier(
        n_estimators=200,  # 200个分类器
        max_depth=2,  # 每个分类器的最大深度为2
        learning_rate=0.5  # 学习率为0.5, 即每个分类器的权重缩减系数为0.5
    )

    # 创建一个堆叠分类器
    stacking_clf = StackingClassifier(
        estimators=[('lr', logistic_clf), ('rf', random_forest_clf), ('svc', svm_clf)],
        final_estimator=LogisticRegression()  # 最终分类器选用逻辑回归分类器
    )

    clfs = {
        "logistic_regression": logistic_clf,
        "random_forest": random_forest_clf,
        "svm": svm_clf,
        "hard_voting": voting_clf_hard,
        "soft_voting": voting_clf_soft,
        "bagging": bag_clf,
        "pasting": paste_clf,
        "adaboost": ada_clf,
        "gradient_boosting": xgb_clf,
        "stacking": stacking_clf
    }

    results = []
    best_model = None
    best_accuracy = 0

    for name, clf in clfs.items():
        start_time = time.time()
        clf.fit(X_train, y_train)
        fit_time = time.time() - start_time

        start_time = time.time()
        accuracy = clf.score(X_test, y_test)
        score_time = time.time() - start_time
        logging.info(f"{name} 模型训练完成,用时: {fit_time} 秒")
        logging.info(f"{name} 模型评估完成,用时: {score_time} 秒")

        results.append([name, fit_time, score_time, accuracy])

        # 保存最佳模型
        if accuracy > best_accuracy:
            best_accuracy = accuracy
            best_model = clf

    # 保存最佳模型
    with open("best_model.pkl", "wb") as f:
        pickle.dump(best_model, f)

    headers = ["Classifier", "Training Time (s)", "Prediction Time (s)", "Accuracy"]
    print(tabulate(results, headers=headers, tablefmt="simple"))

if __name__ == '__main__':
    main()

# 基于gradio做模型预测部署

# 加载模型
model = joblib.load('best_model.pkl')

# 定义分类函数
def predict_image(img):
    img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)  # 将图像从BGR格式转换为RGB格式
    img = cv2.resize(img, (32, 32))  # 调整图像大小到32x32
    img = img.reshape(1,-1)  # 将图像展平
    prediction = model.predict(img)
    return "狗" if prediction[0] == 1 else "猫"

# 创建 Gradio 界面
iface = gr.Interface(
    inputs="image",
    outputs="label",
    fn=predict_image
)

# 启动 Gradio 界面
iface.launch()  # 启动界面